import abc
from typing import Generic, TypeVar

import torch

AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')


class MLTrainingEngine(Generic[AnyInput, AnyOutput], abc.ABC):
    @abc.abstractmethod
    def forward(self, input: AnyInput) -> AnyOutput:
        pass

    @abc.abstractmethod
    def backward(self, output: AnyOutput):
        pass

    @abc.abstractmethod
    def optimize(self):
        pass


class MLInferringEngine(Generic[AnyInput, AnyOutput], abc.ABC):
    @abc.abstractmethod
    def forward(self, input: AnyInput) -> AnyOutput:
        pass


class DelegatingInferringEngine(MLInferringEngine[AnyInput, AnyOutput]):
    def __init__(self, engine: MLTrainingEngine[AnyInput, AnyOutput] | MLInferringEngine[AnyInput, AnyOutput]):
        self._training_engine = engine

    def forward(self, input: AnyInput) -> AnyOutput:
        with torch.no_grad():
            return self._training_engine.forward(input)
